import cv2
import numpy as np
import torch
class detect(object):
    def __init__(self, cfg):
        self.cfg =cfg
    def Color_detect(self, image, color):
        '''
        :param image: 输入待检测图像
        :param color: str 颜色类型['red', 'yellow', 'blue']
        :return:
        '''
        img_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        lower_color = np.array(self.cfg[color][0])
        upper_color = np.array(self.cfg[color][1])
        mask = cv2.inRange(img_hsv, lowerb=lower_color, upperb=upper_color)
        cv2.imshow('mask', mask)
        mask = cv2.medianBlur(mask, 15)
        # mask = cv2.morphologyEx(mask,cv2.MORPH_OPEN, (4,4))
        cnts, he = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        maxIndex = 0
        boxes = []
        cls = []
        for i, cnt in enumerate(cnts):
            area = cv2.contourArea(cnt)
            (x, y, w, h) = cv2.boundingRect(cnt)
            boundarea = w*h
            scale_w_h = w / h
            if self.cfg['area'][0] < area < self.cfg['area'][1] and self.cfg['scale_w_h'][0]< scale_w_h < self.cfg['scale_w_h'][1]:
                if float(area/boundarea) > 0.6:
                    maxIndex = i
                    (x, y, w, h) = cv2.boundingRect(cnts[maxIndex])
                    # box = [x, y, x + w, y + h]
                    box = [x-20, y-20, x+w+20, y+h+20]
                    if x+w+20 >= 1280:
                        box[2] = 1280
                    if y+h+20 >= 720:
                        box[3] = 720
                    if x-20 <= 0:
                        box[0] = 10
                    if y-20 <= 0:
                        box[1] = 10
                    boxes.append(box)
                    cls.append(self.cfg[color])
        return boxes, cls